import os

import pandas as pd


class DataWash:
    def __init__(self):
        all_dataset_name = 'SWAT.A1&A2_Dec2015'
        print("start")

        file_name = None
        for file in os.listdir(all_dataset_name):
            if file[-5:] == '1.csv':
                file_name = file
            else:
                continue
        self.df = pd.read_csv(f'{all_dataset_name}/{file_name}')
        self.df['timestamps'] = self.df['Timestamp']
        self.df['timestamps'] = self.df['timestamps'].apply(
            lambda x: pd.Timestamp.timestamp(pd.Timestamp(self.date_wash(x))))
        self.df.drop(columns=['Timestamp'], inplace=True)
        self.get_label()
        self.df.set_index('timestamps', inplace=True)
        self.df.fillna(method='ffill', inplace=True)
        self.df.fillna(method='bfill', inplace=True)
        self.df.dropna(axis=1, how="all", inplace=True)
        file_dir = 'SWAT'
        if file_dir not in os.listdir():
            os.mkdir(file_dir)
        self.df.to_csv(f'{file_dir}/train.csv')

    @staticmethod
    def date_wash(date):
        date = date.split(' ')
        while '' in date:
            date.remove('')
        target = date[0].split('/')
        day, month, year = target[0], target[1], target[2]
        date = f"{year}-{month}-{day} " + " ".join(date[1:])
        return date

    def get_label(self):
        label = {
            'Normal': '1',
            'Attack': '0'
        }
        for key in label:
            self.df.replace(key, label[key], inplace=True)


if __name__ == "__main__":
    app = DataWash()
